package rabbit.flt.plugins.reactor.http.plugin.impl;

import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.HttpResponseStatus;
import org.reactivestreams.Publisher;
import org.springframework.http.HttpMethod;
import org.springframework.http.client.reactive.ClientHttpConnector;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.client.ClientRequest;
import rabbit.flt.common.AbstractConfigFactory;
import rabbit.flt.common.AgentConfig;
import rabbit.flt.common.Headers;
import rabbit.flt.common.trace.MessageType;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.trace.io.HttpRequest;
import rabbit.flt.common.trace.io.HttpResponse;
import rabbit.flt.common.utils.ReflectUtils;
import rabbit.flt.common.utils.StringUtils;
import rabbit.flt.plugins.common.MethodCallback;
import rabbit.flt.plugins.common.plugin.SupportPlugin;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;
import reactor.netty.ByteBufMono;
import reactor.netty.NettyOutbound;
import reactor.netty.http.client.HttpClient;
import reactor.netty.http.client.HttpClientConfig;
import reactor.netty.http.client.HttpClientRequest;
import reactor.netty.http.client.HttpClientResponse;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.net.URI;
import java.util.HashMap;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;

import static rabbit.flt.common.utils.ReflectUtils.callMethod;
import static rabbit.flt.plugins.reactor.http.plugin.HttpClientFinalizerPlugin.CONTEXT_KEY_NAME;

public class ResponseSinglePlugin extends SupportPlugin {

    /**
     * 代理标记
     */
    protected static ThreadLocal<Boolean> proxyFlagContext = new ThreadLocal<>();

    @Override
    public Object[] before(Object objectEnhanced, Method method, Object[] args) {
        if (isTraceOpened()) {
            args[0] = proxyResponseSingleHandler((BiFunction<HttpClientResponse, ByteBufMono, Mono<?>>) args[0]);
        }
        return args;
    }

    /**
     * 代理对象调用时不做任何trace相关的动作。等同于没有开启trace
     *
     * @return
     */
    @Override
    protected boolean isTraceOpened() {
        return null == proxyFlagContext.get() && super.isTraceOpened();
    }

    /**
     * 代理处理函数
     *
     * @param receiver
     * @return
     */
    private BiFunction<HttpClientResponse, ByteBufMono, Mono<?>> proxyResponseSingleHandler(BiFunction<HttpClientResponse, ByteBufMono, Mono<?>> receiver) {
        return (httpClientResponse, byteBufMono) -> Mono.deferContextual(ctx -> {
            TraceData traceData = ctx.get(CONTEXT_KEY_NAME);
            HttpResponse response = traceData.getHttpResponse();
            httpClientResponse.responseHeaders().forEach(entry -> response.addHeader(entry.getKey(), truncate(entry.getValue())));
            response.setStatusCode(httpClientResponse.status().code());
            traceData.updateCost();
            if (HttpResponseStatus.OK.code() != response.getStatusCode()) {
                return byteBufMono.asString().flatMap(body -> {
                    response.setBody(truncate(body));
                    return receiver.apply(httpClientResponse, ByteBufMono.fromString(Mono.just(body)));
                });
            }
            return receiver.apply(httpClientResponse, byteBufMono);
        });
    }

    /**
     * 调用  responseSingle 方法
     *
     * @param objectEnhanced
     * @param method
     * @param callback
     * @param args
     * @return
     */
    @Override
    public Object invokeMethod(Object objectEnhanced, Method method, MethodCallback callback, Object[] args) {
        if (isTraceOpened()) {
            try {
                TraceData traceData = createTraceData(method);
                proxyFlagContext.set(Boolean.TRUE);
                // 发起tcp请求的时候替换原始对象
                HttpClient httpClient = getHttpClient(objectEnhanced, traceData);
                Object result = callMethod(httpClient, method, args);
                return ((Mono<Object>) result).contextWrite(ctx -> ctx.put(CONTEXT_KEY_NAME, traceData))
                        .doOnSuccess(s -> traceData.updateCost())
                        .onErrorResume(handleError(traceData))
                        .doFinally(s -> handleTraceData(s, traceData));
            } finally {
                proxyFlagContext.remove();
            }
        }
        return super.invokeMethod(objectEnhanced, method, callback, args);
    }

    protected final HttpClient getHttpClient(Object objectEnhanced, TraceData traceData) {
        HttpClient client = (HttpClient) objectEnhanced;
        HttpClientConfig clientConfig = client.configuration();
        Field bodyField = ReflectUtils.loadField(HttpClientConfig.class, "body");
        BiFunction<? super HttpClientRequest, ? super NettyOutbound, ? extends Publisher<Void>> body = ReflectUtils.getValue(clientConfig, bodyField);
        if (null != body) {
            BiFunction<? super HttpClientRequest, ? super NettyOutbound, ? extends Publisher<Void>> proxy = (request, out) -> {
                HttpHeaders httpHeaders = request.requestHeaders();
                Map<String, String> map = clearKeyHeaders(body);
                addTraceHeaders(traceData, clientConfig, httpHeaders);
                Map<String, Object> headers = traceData.getHttpRequest().getHeaders();
                map.forEach(headers::put);
                return body.apply(request, out);
            };
            ReflectUtils.setValue(clientConfig, bodyField, proxy);
            return client;
        } else {
            return client.headers(httpHeaders -> addTraceHeaders(traceData, clientConfig, httpHeaders));
        }
    }

    private Map<String, String> clearKeyHeaders(BiFunction<? super HttpClientRequest, ? super NettyOutbound, ? extends Publisher<Void>> body) {
        for (Field field : body.getClass().getDeclaredFields()) {
            Object value = ReflectUtils.getValue(body, field);
            if (null == value || value instanceof ClientHttpConnector || value instanceof URI || value instanceof HttpMethod) {
                continue;
            }
            for (Field f : value.getClass().getDeclaredFields()) {
                Object sv = ReflectUtils.getValue(value, f);
                if (null != sv && sv instanceof ClientRequest) {
                    return clearKeyHeaders((ClientRequest) sv);
                }
            }
        }
        return new HashMap<>();
    }

    private Map<String, String> clearKeyHeaders(ClientRequest sv) {
        org.springframework.http.HttpHeaders headers = sv.headers();
        MultiValueMap<String, String> headerMap = ReflectUtils.getValue(headers, ReflectUtils.loadField(org.springframework.http.HttpHeaders.class,
                "headers"));
        headerMap.remove(Headers.SPAN_ID);
        headerMap.remove(Headers.TRACE_ID);
        headerMap.remove(Headers.SOURCE_APP);
        return new HashMap<>(headerMap.toSingleValueMap());
    }

    /**
     * 添加trace header信息
     *
     * @param traceData
     * @param clientConfig
     * @param httpHeaders
     */
    protected void addTraceHeaders(TraceData traceData, HttpClientConfig clientConfig, HttpHeaders httpHeaders) {
        if (traceData.isValidTrace()) {
            httpHeaders.set(Headers.SPAN_ID, traceData.getSpanId());
            httpHeaders.set(Headers.TRACE_ID, traceData.getTraceId());
            AgentConfig config = AbstractConfigFactory.getConfig();
            if (null != config && !StringUtils.isEmpty(config.getApplicationCode())) {
                httpHeaders.set(Headers.SOURCE_APP, config.getApplicationCode());
            }
        }
        HttpRequest request = new HttpRequest();
        request.setRequestUri(clientConfig.uri());
        request.setMethod(clientConfig.method().name());
        traceData.setHttpRequest(request);
        httpHeaders.forEach(entry -> request.addHeader(entry.getKey(), truncate(entry.getValue())));
    }

    /**
     * 只设置错误，不发送，finally会发送
     *
     * @param traceData
     * @return
     */
    protected final Function<Throwable, Mono<Object>> handleError(TraceData traceData) {
        return e -> {
            traceData.getHttpResponse().setBody(truncate(e.getMessage()));
            traceData.getHttpResponse().setStatusCode(HttpResponseStatus.INTERNAL_SERVER_ERROR.code());
            traceData.updateCost();
            return Mono.error(e);
        };
    }

    /**
     * 创建trace数据
     *
     * @param method
     * @return
     */
    @Override
    protected final TraceData createTraceData(Method method) {
        TraceData traceData = super.createTraceData(method);
        traceData.setNodeName("ReactorHttp");
        traceData.setMessageType(MessageType.HTTP.name());
        return traceData;
    }

    /**
     * 发送数据
     *
     * @param type
     * @param traceData
     */
    protected void handleTraceData(SignalType type, TraceData traceData) {
        if (type == SignalType.CANCEL && 0 == traceData.getHttpResponse().getStatusCode()) {
            traceData.getHttpResponse().setBody("request is canceled");
        }
        super.handleTraceData(traceData);
    }
}
